import subprocess
import os
from datetime import datetime, timedelta

import psutil
import json

from libs import config
from libs.log import Log


class Scheduler:
    def __init__(self):
        self.logger = Log('scheduler')
        self.running_file_path = os.path.join(os.path.dirname(os.path.abspath(__file__)), '../runtime/tasks.running')
        self.pending_file_path = os.path.join(os.path.dirname(os.path.abspath(__file__)), '../runtime/tasks.pending')
        self.datetime_format = '%Y-%m-%d %H:%M:%S'

    # 检查进程是否仍在运行
    def is_process_running(self, pid):
        return psutil.pid_exists(int(pid))

    # 读取运行中的进程
    def read_running_tasks(self):
        try:
            with open(self.running_file_path, 'r') as f:
                # 使用 json.load() 方法加载数据
                tasks = json.load(f)
                running_tasks = {}
                for pid in tasks:
                    if self.is_process_running(pid):
                        running_tasks[pid] = tasks.get(pid)
        except FileNotFoundError:
            running_tasks = {}
        except json.JSONDecodeError:
            running_tasks = {}
        self.write_running_tasks(running_tasks)
        return running_tasks

    def write_running_tasks(self, tasks):
        with open(self.running_file_path, 'w') as f:
            json.dump(tasks, f, indent=4)

    # 获取任务列表
    def read_tasks(self):
        try:
            with open(self.pending_file_path, 'r') as f:
                # 使用 json.load() 方法加载数据
                tasks = json.load(f)
        except FileNotFoundError:
            tasks = {}
        except json.JSONDecodeError:
            tasks = {}
        return tasks

    def add_task(self, script, user_id, args=[], execute_time: datetime = None, delay_seconds=0, ignore_running=True):
        new_id = self.get_task_id(user_id, script)
        if execute_time is None:
            execute_time = datetime.now()
        if delay_seconds is not None:
            execute_time = execute_time + timedelta(seconds=delay_seconds)
        time_str = execute_time.strftime(self.datetime_format)
        command = f"python {script} {user_id}"
        if len(args) > 0:
            command += f" {' '.join(args)}"
        tasks = self.read_tasks()
        if new_id in tasks:
            return 0
        running_tasks = self.read_running_tasks()
        is_running = False
        for running_task in running_tasks.values():
            if new_id == self.get_task_id(running_task.get('user_id'), running_task.get('command').split(' ')[1]):
                is_running = True
                break
        if not ignore_running and is_running:
            return 0
        self.logger.info('add_task ' + json.dumps({'command': command, 'execute_time': time_str}))
        tasks[new_id] = {'command': command, 'execute_time': time_str, 'create_time': datetime.now().strftime(self.datetime_format)}
        self.save_tasks(tasks)
        return 1

    def get_task_id(self, user_id, script):
        return f"{user_id}_{str(script).replace('.py', '')}"

    def get_channel_from_command(self, cmd):
        script = cmd.split(' ')[1].replace('.py', '')
        if script == 'farm':
            soil_id = cmd.split(' ')[3]
            channel = f"farm_{soil_id}"
        else:
            channel = script
        return channel

    def next_task(self):
        pending_tasks = self.read_tasks()
        running_tasks = self.read_running_tasks()
        process_limit = config.get_process_limit()
        exclude_accounts = config.get_exclude_accounts()
        channel_process_counts = {}
        running_users = []
        for pid in running_tasks:
            running_task = running_tasks.get(pid)
            running_users.append(str(running_task.get('user_id')))
            channel = self.get_channel_from_command(running_task.get('command'), process_limit)
            channel_process_counts[channel] = channel_process_counts.get(channel, 0) + 1
            if channel.startswith('farm_'):
                channel_process_counts['farm'] = channel_process_counts.get('farm', 0) + 1

        sleep_configs = config.get_sleep_configs()
        for task_id in pending_tasks:
            user_id = task_id.split('_')[0]
            # 检查账号是否在休息
            if user_id in sleep_configs:
                sleep_config = sleep_configs.get(user_id)
                start = sleep_config[0]
                end = sleep_config[1]
                hour = datetime.now().hour
                if end > start and start <= hour <= end:
                    continue
                if end < start and (hour >= start or hour <= end):
                    continue

            pending_task = pending_tasks.get(task_id)

            if int(user_id) in exclude_accounts:
                continue
            if str(user_id) in running_users:
                continue
            if pending_task['execute_time'] > datetime.now().strftime(self.datetime_format):
                continue

            channel = self.get_channel_from_command(pending_task.get('command'), process_limit)
            channel_limit = process_limit.get(channel, 1)
            channel_running = channel_process_counts.get(channel, 0)
            if channel_running >= channel_limit:
                continue

            if channel.startswith('farm_'):
                channel = 'farm'
                channel_limit = process_limit.get(channel, 1)
                channel_running = channel_process_counts.get(channel, 0)

                if channel_running >= channel_limit:
                    continue

            pending_task['id'] = task_id
            pending_task['user_id'] = user_id
            return pending_task
        return None


    def save_tasks(self, tasks):
        sorted_tasks = dict(sorted(tasks.items(), key=lambda x: x[1]['execute_time']))
        with open(self.pending_file_path, 'w') as f:
            json.dump(sorted_tasks, f, indent=4)

    def remove_task(self, task_id):
        tasks = self.read_tasks()
        if task_id in tasks:
            tasks.pop(task_id)
            self.save_tasks(tasks)

    def run(self):
        running_tasks = self.read_running_tasks()
        # 读取待执行的任务列表
        task = self.next_task()
        if task is None:
            self.logger.debug(
                f"no_coming_tasks: waiting_count={len(self.read_tasks())}, running_count={len(running_tasks)}")
            return

        self.remove_task(task.get('id'))
        command = task.get('command')
        working_dir = os.path.join(os.path.dirname(os.path.abspath(__file__)), '../')
        env = os.environ.copy()
        log_file_path = os.path.join(os.path.dirname(os.path.abspath(__file__)), f"../runtime/debug/debug_{task.get('user_id')}.log")
        with open(log_file_path, 'a') as f:
            process = subprocess.Popen(
                command,
                shell=True,
                start_new_session=True,
                cwd=working_dir,
                env=env,
                stdout=f,
                stderr=f
            )
        task_pid = process.pid
        self.logger.info('run_task ' + json.dumps({'command': command}))
        running_tasks = self.read_running_tasks()
        running_tasks[task_pid] = {
            "start_time": datetime.now().strftime(self.datetime_format),
            "command": command,
            "user_id": task.get('user_id'),
        }
        self.write_running_tasks(running_tasks)
